# routes/ollama.py

from flask import Blueprint, request, jsonify
from langchain_ollama import ChatOllama
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.prompts import ChatPromptTemplate
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_core.runnables import RunnableLambda, RunnableMap


# 定义蓝图
ollama_bp = Blueprint('ollama', __name__, url_prefix='/ollama')  # 使用 url_prefix 使所有路由都以 /ollama 开头

# 配置 Ollama
OLLAMA_HOST = "http://localhost:11434"
OLLAMA_MODEL = "llama3"  # 可以根据需要调整模型名

# 初始化 Ollama 模型和嵌入
llm = ChatOllama(model=OLLAMA_MODEL)
embeddings = OllamaEmbeddings(model="nomic-embed-text")


@ollama_bp.route('/ask', methods=['POST'])
def ask_ollama():
    data = request.get_json()
    question = data.get("question", "").strip()

    if not question:
        return jsonify({'error': 'Missing or empty question'}), 400

    # 示例文档：可以是来自数据库或其他源的文档
    text = """
    Datawhale 是一个专注于数据科学与 AI 领域的开源组织，汇集了众多领域院校和知名企业的优秀学习者，聚集了一群有开源精神和探索精神的团队成员。
    Datawhale 以"for the learner"为愿景，和学习者一起成长，鼓励真实地展现自我、开放包容、互信互助、敢于试错和勇于担当。
    """

    # 文本分割
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=20)
    chunks = text_splitter.split_text(text)

    # 创建向量存储
    vectorstore = FAISS.from_texts(chunks, embeddings)
    retriever = vectorstore.as_retriever()

    # 创建提示模板
    template = """只能使用下列内容回答问题:
    {context}

    Question: {question}
    """
    prompt = ChatPromptTemplate.from_template(template)

    # 创建检索-问答链
    chain = (
            RunnableMap({
                "context": RunnableLambda(lambda x: "\n".join([
                    doc.page_content for doc in retriever.invoke(x["question"])
                ])),
                "question": RunnableLambda(lambda x: x["question"])
            })
            | prompt
            | llm
    )

    try:
        response = chain.invoke({"question": question})
        answer = response.content if hasattr(response, "content") else str(response)

        return jsonify({
            'question': question,
            'answer': answer,
            'status': 'success'
        })

    except Exception as e:
        return jsonify({'error': str(e)}), 500